import os
import shutil

def cleanup():
    """清理工作目录，保留核心文件"""
    # 要删除的临时文件列表
    temp_files = [
        "test_patch.patch", "result.patch", "final_patch.patch", 
        "ultimate_test.patch", "success.patch", "new_patch.patch",
        "generate_test_patch.py", "demo_patch_editor.py", "redo_demo.py",
        "dead_simple_demo.py", "final_demo.py", "install_correct_patch_lib.py",
        "install_patch_deps.py", "final_final_demo.py", "package_script.py"
    ]
    
    # 要删除的目录
    temp_dirs = ["__pycache__", "build", "dist", "git_patch_editor_gui.spec"]
    
    # 删除文件
    for file in temp_files:
        if os.path.exists(file):
            os.remove(file)
            print(f"🗑️ 删除临时文件: {file}")
    
    # 删除目录
    for dir in temp_dirs:
        if os.path.exists(dir):
            if os.path.isdir(dir):
                shutil.rmtree(dir)
                print(f"🗑️ 删除临时目录: {dir}")
            else:
                os.remove(dir)
                print(f"🗑️ 删除临时文件: {dir}")
    
    # 保留的核心文件
    keep_files = ["git_patch_editor_gui.py", "cleanup_files.py"]
    print("\n✅ 清理完成！保留核心文件:")
    for file in keep_files:
        if os.path.exists(file):
            print(f"   - {file}")

if __name__ == "__main__":
    cleanup()